# -*- coding: utf-8 -*-
"""
功能：
1. find out all the class labels and their numbers in xml files.
   从xml文件中找出所有目标类别以及目标个数

python version: 3.6

usage:

python find_classes_from_xml.py VOC_ROOT_DIR

or

change the default ROOT_DIR in this python script

"""
import os
import sys
import argparse
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


def cal_classes(root_dir):
    xml_dir = os.path.join(root_dir, "Annotations")
    total = len(os.listdir(xml_dir))
    classes_count = {}
    num = 0
    for xml_file in os.listdir(xml_dir):

        # remove 'space' in xml file in field 'filename'
        tree = ET.parse(os.path.join(xml_dir, xml_file))
        root = tree.getroot()
        for obj in root.iter('object'):
            object_class = obj.find('name').text
            if object_class in classes_count:
                classes_count[object_class] += 1
            else:
                classes_count.update({object_class: 1})
        num += 1
        if num % 100 == 0:
            print('processed: {}, total: {}'.format(num, total))
    print(classes_count)


if __name__ == '__main__':

    PARSER = argparse.ArgumentParser(description="arguments")
    PARSER.add_argument("--root_dir", type=str,
                        default="/zongshiban_data/2/ccj/dataset/plane_ship_tank_VOC",
                        help="VOC dataset root directory")
    ARGS = PARSER.parse_args()
    ROOT_DIR = ARGS.root_dir
    cal_classes(ROOT_DIR)
